use tokio::task::JoinHandle;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use crate::base::Task;

use crate::base::TaskHandler;
use crate::persistence::{TaskPersistence, TaskStatus};

#[derive(Debug, thiserror::Error)]
pub enum SpiderError {
    #[error("{0}")]
    Custom(String),
}

pub struct Executor {
    pending_tasks: Mutex<HashMap<String, Arc<Box<dyn Task>>>>,
    running_tasks: Mutex<HashMap<String, Arc<Box<dyn Task>>>>,
    running_task_handlers: Mutex<HashMap<String, Arc<JoinHandle<()>>>>,
    persistence: Arc<dyn TaskPersistence + Send + Sync>,
}

impl Executor {
    pub async fn recover_from_db(persistence: Arc<dyn TaskPersistence + Send + Sync>) -> Result<Self, SpiderError> {
        let executor = Self::new(persistence);
        let persistence = executor.persistence.clone();
        
        // 从持久化存储中获取所有任务
        let tasks = persistence.get_all_tasks()
            .await
            .map_err(|e| SpiderError::Custom(e.to_string()))?;
        
        // 根据任务状态将任务放入对应的HashMap中
        for (id, task, status) in tasks {
            match status {
                TaskStatus::Pending => {
                    executor.pending_tasks
                        .lock()
                        .map_err(|e| SpiderError::Custom(e.to_string()))?
                        .insert(id, task);
                }
                TaskStatus::Running => {
                    executor.start_task2(id, task).await?;
                }
                _ => {}
            }
        }
        
        Ok(executor)
    }

    pub fn new(persistence: Arc<dyn TaskPersistence + Send + Sync>) -> Self {
        Self {
            pending_tasks: Mutex::new(HashMap::new()),
            running_tasks: Mutex::new(HashMap::new()),
            running_task_handlers: Mutex::new(HashMap::new()),
            persistence,
        }
    }

    pub async fn add_task(&self, task: impl Task + 'static) -> Result<TaskHandler, SpiderError> {
        let handler = TaskHandler::new();
        let task = Arc::new(Box::new(task) as Box<dyn Task>);
        self.pending_tasks.lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .insert(handler.id.clone(), task.clone());
        
        // 保存任务到持久化存储
        self.persistence.save_task(&handler.id, &task, TaskStatus::Pending)
            .await
            .map_err(|e| SpiderError::Custom(e.to_string()))?;

        Ok(handler)
    }

    async fn start_task2(&self, id: String, task: Arc<Box<dyn Task>>) -> Result<(), SpiderError> {
        // 使用 tokio 运行时来执行异步任务
        let task_clone = task.clone();
        let persistence = self.persistence.clone();
        let task_id = id.clone();
        let join = tokio::spawn(async move {
            let result = task_clone.run().await;
            match result {
                Ok(_) => {
                    if let Err(e) = persistence.update_task_status(&task_id, TaskStatus::Completed).await {
                        eprintln!("Failed to update task status: {}", e);
                    }
                }
                Err(e) => {
                    eprintln!("Task error: {}", e);
                    if let Err(e) = persistence.update_task_status(&task_id, TaskStatus::Failed).await {
                        eprintln!("Failed to update task status: {}", e);
                    }
                }
            }
        });

        self.running_tasks
        .lock()
        .map_err(|e| SpiderError::Custom(e.to_string()))?
        .insert(id.clone(), task.clone());

        self.running_task_handlers
        .lock()
        .map_err(|e| SpiderError::Custom(e.to_string()))?
        .insert(id.clone(), Arc::new(join));

        Ok(())
    }

    pub async fn start_task(&self, handler: &TaskHandler) -> Result<(), SpiderError> {
        let id = handler.id.clone();
        let task = self.pending_tasks
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .remove(&id)
            .ok_or_else(|| SpiderError::Custom("Task not found".to_string()))?;
        
        // 更新任务状态为运行中
        self.persistence.update_task_status(&id, TaskStatus::Running)
            .await
            .map_err(|e| SpiderError::Custom(e.to_string()))?;

        self.start_task2(id, task).await?;

        Ok(())
    }

    pub async fn start_all_tasks(&self) -> Result<(), SpiderError> {
        // 获取所有待处理任务的handler
        let pending_task_ids: Vec<String> = self.pending_tasks
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .keys()
            .cloned()
            .collect();

        // 为每个任务创建handler并启动
        for id in pending_task_ids {
            let handler = TaskHandler { id };
            self.start_task(&handler).await?;
        }

        Ok(())
    }

    pub async fn stop_task(&self, handler: &TaskHandler) -> Result<(), SpiderError> {
        let id = &handler.id;
        
        // 检查任务是否在运行中
        let task = match self.running_tasks
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .get(id) {
                Some(task) => task.clone(),
                None => {
                    return Ok(());
                }
            };

        // 尝试停止任务
        if let Err(e) = task.stop().await {
            return Err(SpiderError::Custom(format!("Failed to stop task: {}", e)));
        }

        // 获取并取消JoinHandle
        if let Some(handle) = self.running_task_handlers
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .remove(id) 
        {
            handle.abort();
        }

        // 从running_tasks移除并添加到pending_tasks
        if let Some(task) = self.running_tasks
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?
            .remove(id)
        {
            self.pending_tasks
                .lock()
                .map_err(|e| SpiderError::Custom(e.to_string()))?
                .insert(id.clone(), task.clone());

            // 更新任务状态为待处理
            self.persistence.update_task_status(id, TaskStatus::Pending)
                .await
                .map_err(|e| SpiderError::Custom(e.to_string()))?;
        }

        Ok(())
    }

    pub fn is_running(&self, handler: &TaskHandler) -> Result<bool, SpiderError> {
        let running_tasks_guard = self.running_tasks
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?;

        let running_handlers_guard = self.running_task_handlers
            .lock()
            .map_err(|e| SpiderError::Custom(e.to_string()))?;

        // 检查任务是否在running_tasks中，且对应的JoinHandle未完成
        if let Some(_) = running_tasks_guard.get(&handler.id) {
            if let Some(handle) = running_handlers_guard.get(&handler.id) {
                return Ok(!handle.is_finished());
            }
        }

        Ok(false)
    }
}




#[cfg(test)]
mod tests {

    use std::time::Duration;

    use tokio::time::sleep;
    use crate::{base::{Runnable, Stopable, Task, TaskError}, persistence::{PersistenceError, SeaOrmTaskPersistence}};
    use async_trait::async_trait;

    use super::*;


    pub struct DummyTaskPersistence;

    #[async_trait]
    impl TaskPersistence for DummyTaskPersistence {
        async fn save_task(&self, task_id: &str, _: &Arc<Box<dyn Task>>, status: TaskStatus) -> Result<(), PersistenceError> {
            println!("Saving task {} with status {:?}", task_id, status);
            Ok(())
        }

        async fn remove_task(&self, task_id: &str) -> Result<(), PersistenceError> {
            println!("Removing task {}", task_id);
            Ok(())
        }

        async fn update_task_status(&self, task_id: &str, status: TaskStatus) -> Result<(), PersistenceError> {
            println!("Updating task {} status to {:?}", task_id, status);
            Ok(())
        }

        async fn get_all_tasks(&self) -> Result<Vec<(String, Arc<Box<dyn Task>>, TaskStatus)>, PersistenceError> {
            println!("Getting all tasks");
            Ok(Vec::new())
        }
    }
    

    // 测试stop报错的情况
    #[tokio::test]
    async fn test_simple_task() {
        pub struct SimpleTask1 {}

        impl SimpleTask1 {
            pub fn new() -> Self {
                Self {
                }
            }
        }
        impl Task for SimpleTask1 {
            fn type_name(&self) -> &'static str {
                "SimpleTask1"
            }
            fn serialize(&self) -> Result<String, TaskError> {
                Ok("".to_string())
            }
        }
        #[async_trait]
        impl Runnable for SimpleTask1 {
            async fn run(&self) -> Result<(), TaskError> {
                sleep(Duration::from_secs(5)).await;
                Ok(())
            }
        }
        #[async_trait]
        impl Stopable for SimpleTask1 {
            async fn stop(&self) -> Result<(), TaskError> {
                Err(TaskError::Custom("stop failed".to_string()))
            }
        }

        let task = SimpleTask1::new();
        let executor = Executor::new(Arc::new(DummyTaskPersistence{}));
        let handler = executor.add_task(task).await.unwrap();
        executor.start_task(&handler).await.unwrap();
        assert_eq!(executor.is_running(&handler).unwrap(), true);
        let stop_result = executor.stop_task(&handler).await;
        assert_eq!(stop_result.is_err(), true);
        assert_eq!(executor.is_running(&handler).unwrap(), true);
        sleep(Duration::from_secs(10)).await;
        assert_eq!(executor.is_running(&handler).unwrap(), false);
    }

    // 测试任务暂停后可以继续
    #[tokio::test]
    async fn test_simple_task2() {
        struct SimpleTask1 {
            pub sleep_time: Mutex<i32>,
        }

        impl SimpleTask1 {
            pub fn new() -> Self {
                Self {
                    sleep_time: Mutex::new(10),
                }
            }
        }
        impl Task for SimpleTask1 {
            fn type_name(&self) -> &'static str {
                "SimpleTask1"
            }
            fn serialize(&self) -> Result<String, TaskError> {
                Ok("".to_string())
            }
        }
        #[async_trait]
        impl Runnable for SimpleTask1 {
            async fn run(&self) -> Result<(), TaskError> {
                // 每睡一秒，sleep_time减一
                while *self.sleep_time.lock().unwrap() > 0 {
                    sleep(Duration::from_secs(1)).await;
                    let mut sleep_time = self.sleep_time.lock().unwrap();
                    println!("left time: {}", sleep_time);
                    *sleep_time -= 1; 
                }
                Ok(())
            }
        }
        #[async_trait]
        impl Stopable for SimpleTask1 {
            async fn stop(&self) -> Result<(), TaskError> {
                Ok(())
            }
        }

        let task = SimpleTask1::new();
        let executor = Executor::new(Arc::new(DummyTaskPersistence{}));
        let handler = executor.add_task(task).await.unwrap();
        executor.start_task(&handler).await.unwrap();
        assert_eq!(executor.is_running(&handler).unwrap(), true);
        let stop_result = executor.stop_task(&handler).await;
        assert_eq!(stop_result.is_err(), false);
        assert_eq!(executor.is_running(&handler).unwrap(), false);
        executor.start_task(&handler).await.unwrap();
        assert_eq!(executor.is_running(&handler).unwrap(), true);
        sleep(Duration::from_secs(5)).await;
        assert_eq!(executor.is_running(&handler).unwrap(), true);
        executor.stop_task(&handler).await.unwrap();
        executor.start_task(&handler).await.unwrap();
        sleep(Duration::from_secs(8)).await;
        assert_eq!(executor.is_running(&handler).unwrap(), false);
    }


    // 测试Executor，将任务保存到数据库
    struct DbTask {
        pub sleep_time: Mutex<i32>,
    }

    impl DbTask {
        pub fn new(time: i32) -> Self {
            Self {
                sleep_time: Mutex::new(time),
            }
        }
    }
    impl Task for DbTask {
        fn type_name(&self) -> &'static str {
            "DbTask"
        }
        fn serialize(&self) -> Result<String, TaskError> {
            Ok(self.sleep_time.lock().unwrap().to_string())
        }
    }
    #[async_trait]
    impl Runnable for DbTask {
        async fn run(&self) -> Result<(), TaskError> {
            // 每睡一秒，sleep_time减一
            while *self.sleep_time.lock().unwrap() > 0 {
                sleep(Duration::from_secs(1)).await;
                let mut sleep_time = self.sleep_time.lock().unwrap();
                println!("left time: {}", sleep_time);
                *sleep_time -= 1; 
            }
            Ok(())
        }
    }
    #[async_trait]
    impl Stopable for DbTask {
        async fn stop(&self) -> Result<(), TaskError> {
            Ok(())
        }
    }

    // 测试Executor，将任务保存到数据库
    #[tokio::test]
    async fn test_simple_task3() {
        std::env::set_var("DATABASE_URL", "sqlite://../../db.sqlite");
        let task = DbTask::new(10);
        let persistence = SeaOrmTaskPersistence::new(|_, task_data| {
            let sleep_time = task_data.parse::<i32>().unwrap();
            Ok(Box::new(DbTask::new(sleep_time)) as Box<dyn Task>) 
        }).await;
        let executor = Executor::new(Arc::new(persistence));
        executor.add_task(task).await.unwrap();
    }

    // 测试Executor，从数据库恢复任务
    #[tokio::test]
    async fn test_simple_task4() {
        std::env::set_var("DATABASE_URL", "sqlite://../../db.sqlite");
        let persistence = SeaOrmTaskPersistence::new(|_, task_data| {
            let sleep_time = task_data.parse::<i32>().unwrap();
            Ok(Box::new(DbTask::new(sleep_time)) as Box<dyn Task>) 
        }).await;
        let executor = Executor::recover_from_db(Arc::new(persistence)).await.unwrap();
        executor.start_all_tasks().await.unwrap();

        sleep(Duration::from_secs(15)).await;
    }
}
